package ai.zixing.mashibing.basic_class.class10;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集：结构
 */
public class Code01_UnionFind {

    // 构建单项列表
    public static class Node<V> {
        V value;

        public Node(V value) {
            this.value = value;
        }
    }

    public static class UnionSet<V> {
        // 数据对于表，找到对应的节点，一一对应
        public HashMap<V, Node<V>> nodes;

        // 往上找祖先
        public HashMap<Node<V>, Node<V>> parents;

        // 只有一个点，它是祖先节点，才会有值
        public HashMap<Node<V>, Integer> sizeMap;

        /**
         * 将所有样本初始化
         *
         * @param values 所有样本集合
         */
        public UnionSet(List<V> values) {
            for (V value : values) {
                Node<V> node = new Node<>(value);
                nodes.put(value, node);
                parents.put(node, node);
                // 每个数据都是一个代表点【祖先】
                sizeMap.put(node, 1);
            }
        }

        /**
         * 找到 cur 的代表点【祖先】
         * 树的扁平化优化
         *
         * @param cur 当前节点
         * @return Node<V>
         */
        public Node<V> findFather(Node<V> cur) {
            // 保存找的代表点的路径
            Stack<Node<V>> path = new Stack<>();
            // 找代表点
            while (cur != parents.get(cur)) {
                path.push(cur);
                cur = parents.get(cur);
            }
            // cur 是头节点
            // 优化，路径上所有节点的终极代表点实际就是 cur
            while (!path.isEmpty()) {
                parents.put(path.pop(), cur);
            }
            return cur;
        }

        /**
         * 是否属于同一集合
         *
         * @param a 元素 a
         * @param b 元素 b
         * @return 属于同一个集合 返回 true
         */
        public boolean isSameSet(V a, V b) {
            if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
                return false;
            }
            // 同一个祖先
            return findFather(nodes.get(a)) == findFather(nodes.get(b));
        }

        /**
         * @return 合并之后的数量
         */
        public int getSetNum() {
            return sizeMap.size();
        }


        /**
         * 并集
         *
         * @param a a集合
         * @param b b集合
         */
        public void union(V a, V b) {
            if (!nodes.containsKey(a) || !nodes.containsKey(b)) {
                return;
            }
            // 找到两个集合的代表点
            Node<V> aHead = findFather(nodes.get(a));
            Node<V> bHead = findFather(nodes.get(b));
            // 代表点不相同的时候才合并
            if (aHead != bHead) {
                int aSetSize = sizeMap.get(aHead);
                int bSetSize = sizeMap.get(bHead);
                // 将小集合合并到大集合
                // 尽可能的使合并完集合链表短一些
                if (aSetSize >= bSetSize) {
                    // 小集合的代表点改为大集合的代表点
                    parents.put(bHead, aHead);
                    // 集合大小为两个集合的大小和
                    sizeMap.put(aHead, aSetSize + bSetSize);
                    // 小集合不在作为代表点，删除
                    sizeMap.remove(bHead);
                } else {
                    parents.put(aHead, bHead);
                    sizeMap.put(bHead, aSetSize + bSetSize);
                    sizeMap.remove(aHead);
                }
            }
        }
    }
}
